"""
author: liuxu
date: 2025/9/20
description: 测试$ROCM_PATH/lib/mpitracer.so库功能
"""
import re
import allure
import pytest
import os
from common.logger_ctrl import mylogger
from common.ssh_command import exec_cmd
from setting import CPP_DIR, EXE_DIR
from base_public.testbase import BasePublic
from __check_mpi_available import check_mpi_environment

@pytest.mark.medium
@allure.feature('测试mpitracer.so库功能')
class TestMpiTracer(BasePublic):
    """测试MPI Tracer功能的测试类"""

    @classmethod
    def setup_class(cls):
        """类级别设置，检查MPI环境是否可用"""
        cls.mpi_available, mpi_info = check_mpi_environment()
        if not cls.mpi_available:
            pytest.skip("The MPI environment is unavailable")

        cls.mpi_so_abspath = mpi_info.get("mpi_so_abspath", "libmpi.so")

        # 获取ROCM路径
        rocm_path = os.environ.get("ROCM_PATH", "/opt/dtk")
        cls.mpitracer_so = os.path.join(rocm_path, "lib", "mpitracer.so")

        if not os.path.exists(cls.mpitracer_so):
            mylogger.warning(f"MPI Tracer库未找到: {cls.mpitracer_so}")
            pytest.skip("MPI Tracer库未找到，跳过测试")

        mylogger.info(f"MPI环境检查通过，MPI Tracer库: {cls.mpitracer_so}")

    def test_mpitracer_c_fortran_mixed(self, test_env):
        """
        使用mpitracer拦截C和Fortran混合MPI程序
        """
        #TODO后续考虑self.compile_source_code支持编译多个源码文件
        # 编译测试程序
        exe_path = f"{EXE_DIR}/mpitracer/mixed_c_fortran"
        self.create_directory(os.path.dirname(exe_path))
        res, oouput = exec_cmd(
            f"mpifort {CPP_DIR}/Mpi_case/c_fortran/mixed_c_fortran.f90 {CPP_DIR}/Mpi_case/c_fortran/mixed_c_fortran.c -o {exe_path}")
        assert res is True
        # exe_path = self.compile_source_code(src_file_name="25.04.2/mpitracer_test/mpi_tests_openmpi/mixed_c_fortran.f90",
        #                                     compiler="mpifort")

        # 执行测试命令
        ld_preload = f"{self.mpitracer_so}:{self.mpi_so_abspath}"  # LD_PRELOAD路径
        run_cmd = f"mpirun --allow-run-as-root -x LD_PRELOAD={ld_preload} -np 4 {exe_path}"
        res, output = exec_cmd(run_cmd, cwd=test_env)

        # 校验测试结果正确性
        assert res is True, "The mpi program failed to run"
        # 检查MPI Tracer输出
        assert "MPI_Init" in output, "MPI Tracer未检测到MPI_Init调用"
        assert "MPI_Finalize" in output, "MPI Tracer未检测到MPI_Finalize调用"

